//!
//! no_std环境，适用于运行环境已经有libcrypto.so提供HMAC和BASE64算法.
//! TOTP也经常使用BASE32算法，采用了支持no_std的crate: data_encoding.
//!
#![no_std]

/// TOTP缺省使用HMAC-SHA1,但也可以使用HMAC-SHA256,HMAC-SHA512
#[derive(Debug, Clone, Copy)]
pub enum Algorithm {
    Sha1 = 0,
    Sha256,
    Sha512
}

/// 密钥可以是BASE32/BASE64编码，Raw代表未做任何编码.
#[derive(Debug, Clone, Copy)]
pub enum CodeType {
    Raw = 0,
    Base32,
    Base64
}

/// 提供TOTP算法
pub struct Totp<'a> {
    key: &'a [u8],
    algo: Algorithm,
    code: CodeType,
    digits: u32,
    period: u32,
    t0: i64,
}

impl<'a> Totp<'a> {
    pub fn new(key: &'a [u8], code: CodeType) -> Self {
        Self {
            key,
            code,
            algo: Algorithm::Sha1,
            digits: 6,
            period: 30,
            t0: 0,
        }
    }

    pub fn from_raw(key: &'a [u8]) -> Self {
        Self::new(key, CodeType::Raw)
    }

    pub fn from_base32(key: &'a str) -> Self {
        Self::new(key.as_bytes(), CodeType::Base32)
    }

    pub fn from_base64(key: &'a str) -> Self {
        Self::new(key.as_bytes(), CodeType::Base64)
    }

    /// 缺省值为CodeType::Base32
    /// 如果解码失败，则按照Raw处理.
    pub fn set_code(&mut self, code: CodeType) -> &mut Self {
        self.code = code;
        self
    }

    pub fn get_code(&self) -> CodeType {
        self.code
    }

    /// 缺省值为Algorithm::Sha1
    pub fn set_algorithm(&mut self, algo: Algorithm) -> &mut Self {
        self.algo = algo;
        self
    }

    pub fn get_algorithm(&self) -> Algorithm {
        self.algo
    }

    /// 如果是0，表示缺省值6，如果大于10则截断为10.
    pub fn set_digits(&mut self, digits: u32) -> &mut Self {
        self.digits = match digits {
            0 => 6,
            1..=10 => digits,
            _ => 10,
        };
        self
    }

    pub fn get_digits(&self) -> u32 {
        self.digits
    }

    /// 如果是0，表示缺省值30秒.
    pub fn set_period(&mut self, period: u32) -> &mut Self {
        self.period = match period {
            0 => 30,
            _ => period,
        };
        self
    }

    pub fn get_period(&self) -> u32 {
        self.period
    }

    /// 设定时间原点, 缺省为0.
    pub fn set_t0(&mut self, t0: i64) -> &mut Self {
        self.t0 = t0;
        self
    }

    pub fn get_t0(&self) -> i64 {
        self.t0
    }

    /// 基于当前时间计算
    pub fn hash_now(&self) -> u32 {
        let now = Self::now();
        self.hash(now)
    }

    /// 基于给定时间计算，验证的时候需要考虑网络延时，可以利用这个接口进行多个时间区间的验证.
    pub fn hash(&self, time: i64) -> u32 {
        let tm = (time - self.t0) / self.period as i64;
        let mut key = KeyBuf::new(self.key.len());
        let key = self.decode(&mut key);
        match self.algo {
            Algorithm::Sha1 => self.hmac_shaxxx::<_, 20>(key, tm, || unsafe { EVP_sha1() }),
            Algorithm::Sha256 => self.hmac_shaxxx::<_, 32>(key, tm, || unsafe { EVP_sha256() }),
            Algorithm::Sha512 => self.hmac_shaxxx::<_, 64>(key, tm, || unsafe { EVP_sha512() }),
        }
    }

    /// 获取当前UNIX时间.
    pub fn now() -> i64 {
        unsafe { time(core::ptr::null_mut()) }
    }

    fn decode<'b, 'c>(&self, buf: &'b mut KeyBuf) -> &'c [u8] where 'a: 'c, 'b: 'c {
        match self.code {
            CodeType::Raw => self.key,
            CodeType::Base32 => self.decode_base32(buf.buf()),
            CodeType::Base64 => self.decode_base64(buf.buf()),
        }
    }

    fn decode_base32<'b, 'c>(&self, buf: &'b mut [u8]) -> &'c [u8] where 'a: 'c, 'b: 'c {
        let Ok(len) = data_encoding::BASE32.decode_len(self.key.len()) else { return self.key; };
        let buf = &mut buf[0..len];
        match data_encoding::BASE32.decode_mut(self.key, buf) {
            Ok(size) => &buf[..size],
            _ => self.key,
        }
    }

    // libcrypto.so支持base64的解码，尽可能使用，减少rust的二进制大小
    fn decode_base64<'b, 'c>(&self, buf: &'b mut [u8]) -> &'c [u8] where 'a: 'c, 'b: 'c {
        let len = unsafe { EVP_DecodeBlock(buf as *mut _ as *mut u8, self.key as *const _ as *const u8, self.key.len() as i32) };
        if len == -1 {
            return self.key;
        }
        &buf[..len as usize]
    }

    fn hmac_shaxxx<F, const N: usize>(&self, key: &[u8], time: i64, f: F) -> u32 where F: Fn() -> *const EVP_MD {
        let tm = time.to_be_bytes();
        let mut md = [0_u8; N];
        let mut len = N as u32;
        let _ = unsafe { HMAC(f(), key.as_ptr(), key.len() as i32, tm.as_ptr(), 8, md.as_mut_ptr(), &mut len) };
        self.trunc(&md)
    }

    fn trunc(&self, md: &[u8]) -> u32 {
        let offset = (md[md.len() - 1] & 0xF) as usize;
        let val = u32::from_be_bytes([md[offset] & 0x7F, md[offset + 1], md[offset + 2], md[offset + 3]]);
        (val as u64 % 10_u64.pow(self.digits)) as u32
    }
}

const BUFSIZE: usize = 64;

struct KeyBuf {
    key: [u8; BUFSIZE],
    buf: *mut u8,
    size: usize,
}

impl KeyBuf {
    fn new(size: usize) -> Self {
        let buf = if size <= BUFSIZE {
            core::ptr::null_mut()
        } else {
            unsafe { malloc(size) }
        };
        Self {
            key: [0_u8; BUFSIZE],
            buf,
            size,
        }
    }

    fn buf(&mut self) -> &mut[u8] {
        if self.buf.is_null() {
            &mut self.key[0..self.size]
        } else {
            unsafe { core::slice::from_raw_parts_mut(self.buf, self.size) }
        }
    }
}

impl Drop for KeyBuf {
    fn drop(&mut self) {
        if !self.buf.is_null() {
            unsafe { free(self.buf) };
        }
    }
}

#[link(name = "c")]
extern "C" {
    fn time(tm: *mut i64) -> i64;
    fn malloc(size: usize) -> *mut u8;
    fn free(buf: *mut u8);
}

#[repr(C)]
#[allow(non_camel_case_types)]
struct EVP_MD (u8);

#[link(name = "crypto")]
extern "C" {
    fn EVP_DecodeBlock(buf: *mut u8, input: *const u8, len: i32) -> i32;
    fn EVP_sha1() -> *const EVP_MD;
    fn EVP_sha256() -> *const EVP_MD;
    fn EVP_sha512() -> *const EVP_MD;
    fn HMAC(evp_md: *const EVP_MD, key: *const u8, key_len: i32, data: *const u8, data_len: usize, md: *mut u8, md_len: *mut u32) -> *const u8;
}
